# websocket_manager.py
import asyncio
import json
from typing import Dict, List, Optional
from fastapi import WebSocket
import uuid
from utils.log_util import logger


class ConnectionManager:
    def __init__(self):
        # 管理 WebSocket 连接
        self.active_connections: Dict[str, WebSocket] = {}
        # 管理客户端与 WebSocket 的映射
        self.client_to_websocket: Dict[str, List[str]] = {}  # client_id -> [websocket_ids]
        # ✅ 管理 WebSocket 与 user_id 的映射
        self.websocket_user_map: Dict[str, Optional[int]] = {}  # websocket_id -> user_id

    async def connect(self, websocket: WebSocket, client_id: str = None, user_id: Optional[int] = None):
        """WebSocket 连接
        
        Args:
            websocket: WebSocket连接
            client_id: 客户端ID（可选）
            user_id: 用户ID（可选，用于消息过滤）
        """
        await websocket.accept()
        websocket_id = str(uuid.uuid4())
        self.active_connections[websocket_id] = websocket

        if client_id:
            if client_id not in self.client_to_websocket:
                self.client_to_websocket[client_id] = []
            self.client_to_websocket[client_id].append(websocket_id)
        
        # ✅ 保存user_id映射
        if user_id is not None:
            self.websocket_user_map[websocket_id] = user_id
            logger.debug(f"✅ WebSocket user_id映射已保存: websocket_id={websocket_id}, user_id={user_id}, websocket_user_map={dict(self.websocket_user_map)}")
        else:
            logger.warning(f"⚠️ WebSocket连接时user_id为None: websocket_id={websocket_id}, 当前websocket_user_map={dict(self.websocket_user_map)}")

        return websocket_id

    def disconnect(self, websocket_id: str):
        """WebSocket 断开连接"""
        if websocket_id in self.active_connections:
            del self.active_connections[websocket_id]

        # 清理映射关系
        for client_id, websocket_ids in self.client_to_websocket.items():
            if websocket_id in websocket_ids:
                websocket_ids.remove(websocket_id)
            if not websocket_ids:
                del self.client_to_websocket[client_id]
        
        # ✅ 清理user_id映射
        if websocket_id in self.websocket_user_map:
            del self.websocket_user_map[websocket_id]

    async def send_personal_message(self, message: dict, websocket_id: str):
        """向特定 WebSocket 发送消息"""
        if websocket_id in self.active_connections:
            websocket = self.active_connections[websocket_id]
            await websocket.send_json(message)

    async def broadcast_to_client(self, message: dict, client_id: str):
        """向关注特定客户端的 WebSocket 广播消息"""
        if client_id in self.client_to_websocket:
            for websocket_id in self.client_to_websocket[client_id]:
                await self.send_personal_message(message, websocket_id)

    async def broadcast(self, message: dict, user_id: Optional[int] = None):
        """
        向所有 WebSocket 广播（支持按user_id过滤）
        
        Args:
            message: 要广播的消息
            user_id: 如果提供，只推送给该用户的WebSocket连接
        """
        if user_id is not None:
            # ✅ 只推送给指定用户的WebSocket连接
            matched_count = 0
            for ws_id, ws in self.active_connections.items():
                if self.websocket_user_map.get(ws_id) == user_id:
                    try:
                        await ws.send_json(message)
                        matched_count += 1
                    except Exception as e:
                        logger.warning(f"WebSocket推送失败: ws_id={ws_id}, user_id={user_id}, error={e}")
            
            if matched_count == 0:
                # ⚠️ 警告：指定user_id但没有匹配的连接
                logger.warning(
                    f"WebSocket广播：user_id={user_id}没有匹配的连接。"
                    f"当前连接数={len(self.active_connections)}, "
                    f"websocket_user_map={dict(self.websocket_user_map)}"
                )
        else:
            # 某些消息类型应该广播给所有用户（如status_update、client_update等）
            message_type = message.get('type')
            event = message.get('event')
            # 修复：get_address事件需要广播给所有用户（管理员需要查看agent的地址列表）
            should_broadcast_to_all = (
                message_type in ('status_update', 'client_update', 'client_disconnected', 'init') or
                (message_type == 'client_response' and event == 'get_address')
            )
            
            if should_broadcast_to_all:
                # 这些消息类型应该广播给所有用户，不需要警告
                logger.debug(f"WebSocket广播：{message_type}消息广播给所有用户（正常行为），当前连接数={len(self.active_connections)}")
            else:
                # 警告：业务消息（如search、order等）user_id为None，可能导致消息泄露
                logger.warning(
                    f"⚠️ WebSocket广播：user_id为None，广播给所有连接（可能导致消息泄露）。"
                    f"消息类型={message_type}, 当前连接数={len(self.active_connections)}"
                )
            
            # 广播给所有连接
            for websocket in self.active_connections.values():
                try:
                    await websocket.send_json(message)
                except Exception as e:
                    logger.warning(f"WebSocket推送失败: error={e}")


# 全局 WebSocket 管理器
websocket_manager = ConnectionManager()